
#include <opencv2/opencv.hpp>
#include <opencv2/highgui/highgui.hpp>
#include <opencv2/imgproc/imgproc.hpp>
#include "onnxruntime.h"

onnxruntime::onnxruntime(std::wstring model_path, int num_threads = 1, std::vector<int64_t> input_node_dims = { 1, 3, 128, 128 }) {
	input_node_dims_ = input_node_dims;
	for (int64_t i : input_node_dims_) {
		input_tensor_size_ *= i;
		out_tensor_size_ *= i;
	}

	//std::cout << input_tensor_size_ << std::endl;
	session_options_.SetIntraOpNumThreads(num_threads);
	session_options_.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);

	try {
		session_ = Ort::Session(env_, model_path.c_str(), session_options_);
	}
	catch (...) {

	}

	Ort::AllocatorWithDefaultOptions allocator;
	//获取输入name
	const char* input_name = session_.GetInputName(0, allocator);
	input_node_names_ = { input_name };
	//std::cout << "input name:" << input_name << std::endl;
	const char* output_name = session_.GetOutputName(0, allocator);
	out_node_names_ = { output_name };
	//std::cout << "output name:" << output_name << std::endl;
	input_size = cv::Size(int(input_node_dims_[3]), int(input_node_dims_[2]));
}

void onnxruntime::predict(cv::Mat& blob)
{
	clock_t start{ clock() }, end;
	auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);

	input_tensors_.emplace_back(Ort::Value::CreateTensor<float>(memory_info, blob.ptr<float>(), blob.total(), input_node_dims_.data(), input_node_dims_.size()));

	std::vector<Ort::Value> output_tensors_ = session_.Run(
		Ort::RunOptions{ nullptr },
		input_node_names_.data(),
		input_tensors_.data(),
		input_node_names_.size(),
		out_node_names_.data(),
		out_node_names_.size());
	float* floatarr = output_tensors_[0].GetTensorMutableData<float>();
	end = clock();
	std::cout << "onnxruntime:" << end - start << "ms" << std::endl;
	for (int i = 0; i < 4; i++) {
		std::cout << *floatarr++ << std::endl;
	}
}


